package com.zc.ai.demo;

import com.theokanning.openai.completion.chat.*;
import com.theokanning.openai.completion.CompletionRequest;
import com.theokanning.openai.image.CreateImageRequest;
import com.theokanning.openai.model.Model;
import com.zc.ai.service.openai.OpenAiService;
import io.reactivex.Flowable;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;

class OpenAiApiExample {
    public static void main(String... args) {
        System.setProperty("java.net.useSystemProxies", "true");
        System.setProperty("https.proxyHost", "127.0.0.1");
        System.setProperty("https.proxyPort", "7890");
//        String token = System.getenv("OPENAI_TOKEN");
        String token = "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";
        OpenAiService service = new OpenAiService(token);
        List<Model> models = service.listModels();
        System.out.println(models);

//        System.out.println("\nCreating completion...");
        System.out.println("\n正在完成创建...");
        CompletionRequest completionRequest = CompletionRequest.builder()
                /**
                 * 模型选择 ada
                 * */
                .model("ada")
                /**
                 * 提示  [有人曾经告诉我，世界将把我卷起来]
                 * */
//                .prompt("Somebody once told me the world is gonna roll me")
                .prompt("有人曾经告诉我，世界将把我卷起来")
                .echo(true)
                /**
                 * 用户 [用户]
                 * */
                .user("testing")
                .n(3)
                .build();
        service.createCompletion(completionRequest).getChoices().forEach(System.out::println);

//        System.out.println("\nCreating Image...");
        System.out.println("\n正在创建图像...");
        CreateImageRequest request = CreateImageRequest.builder()
                /**
                 * 提示 [一头跳霹雳舞的牛]
                 * */
                .prompt("A cow breakdancing with a turtle")
                .build();

                /**
                 * 图像位于：service.createImage(request).getData().get(0).getUrl()
                 * */
        System.out.println("\nImage is located at:");
        System.out.println(service.createImage(request).getData().get(0).getUrl());

        System.out.println("Streaming chat completion...");
        final List<ChatMessage> messages = new ArrayList<>();
        final ChatMessage systemMessage = new ChatMessage(ChatMessageRole.SYSTEM.value(), "你是一条狗，会这样说话");
        messages.add(systemMessage);
        ChatCompletionRequest chatCompletionRequest = ChatCompletionRequest
                .builder()
                .model("gpt-3.5-turbo")
                .messages(messages)
                .n(1)
                .maxTokens(50)
                .logitBias(new HashMap<>())
                .build();

        Flowable<ChatCompletionChunk> chatCompletionChunkFlowable = service.streamChatCompletion(chatCompletionRequest);
        service.streamChatCompletion(chatCompletionRequest)
                .doOnError(Throwable::printStackTrace)
                .blockingForEach(System.out::println);

        service.shutdownExecutor();
    }
}
